{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hjM_sV_AepYf"
   },
   "source": [
    "# 🔪 JAX - The Sharp Bits 🔪\n",
    "\n",
    "<!--* freshness: { reviewed: '2024-06-03' } *-->\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jax-ml/jax/blob/main/docs/notebooks/Common_Gotchas_in_JAX.ipynb) [![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/jax-ml/jax/blob/main/docs/notebooks/Common_Gotchas_in_JAX.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4k5PVzEo2uJO"
   },
   "source": [
    "When walking about the countryside of Italy, the people will not hesitate to tell you that __JAX__ has [_\"una anima di pura programmazione funzionale\"_](https://www.sscardapane.it/iaml-backup/jax-intro/).\n",
    "\n",
    "__JAX__ is a language for __expressing__ and __composing__ __transformations__ of numerical programs. __JAX__ is also able to __compile__ numerical programs for CPU or accelerators (GPU/TPU).\n",
    "JAX works great for many numerical and scientific programs, but __only if they are written with certain constraints__ that we describe below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "GoK_PCxPeYcy"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from jax import jit\n",
    "from jax import lax\n",
    "from jax import random\n",
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gX8CZU1g2agP"
   },
   "source": [
    "## 🔪 Pure functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2oHigBkW2dPT"
   },
   "source": [
    "JAX transformation and compilation are designed to work only on Python functions that are functionally pure: all the input data is passed through the function parameters, all the results are output through the function results. A pure function will always return the same result if invoked with the same inputs.\n",
    "\n",
    "Here are some examples of functions that are not functionally pure for which JAX behaves differently than the Python interpreter. Note that these behaviors are not guaranteed by the JAX system; the proper way to use JAX is to use it only on functionally pure Python functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "A6R-pdcm4u3v",
    "outputId": "25dcb191-14d4-4620-bcb2-00492d2f24e1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing function\n",
      "First call:  4.0\n",
      "Second call:  5.0\n",
      "Executing function\n",
      "Third call, different type:  [5.]\n"
     ]
    }
   ],
   "source": [
    "def impure_print_side_effect(x):\n",
    "  print(\"Executing function\")  # This is a side-effect\n",
    "  return x\n",
    "\n",
    "# The side-effects appear during the first run\n",
    "print (\"First call: \", jit(impure_print_side_effect)(4.))\n",
    "\n",
    "# Subsequent runs with parameters of same type and shape may not show the side-effect\n",
    "# This is because JAX now invokes a cached compilation of the function\n",
    "print (\"Second call: \", jit(impure_print_side_effect)(5.))\n",
    "\n",
    "# JAX re-runs the Python function when the type or shape of the argument changes\n",
    "print (\"Third call, different type: \", jit(impure_print_side_effect)(jnp.array([5.])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "-N8GhitI2bhD",
    "outputId": "fd3624c9-197d-42cb-d97f-c5e0ef885467"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First call:  4.0\n",
      "Second call:  5.0\n",
      "Third call, different type:  [14.]\n"
     ]
    }
   ],
   "source": [
    "g = 0.\n",
    "def impure_uses_globals(x):\n",
    "  return x + g\n",
    "\n",
    "# JAX captures the value of the global during the first run\n",
    "print (\"First call: \", jit(impure_uses_globals)(4.))\n",
    "g = 10.  # Update the global\n",
    "\n",
    "# Subsequent runs may silently use the cached value of the globals\n",
    "print (\"Second call: \", jit(impure_uses_globals)(5.))\n",
    "\n",
    "# JAX re-runs the Python function when the type or shape of the argument changes\n",
    "# This will end up reading the latest value of the global\n",
    "print (\"Third call, different type: \", jit(impure_uses_globals)(jnp.array([4.])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "RTB6iFgu4DL6",
    "outputId": "16697bcd-3623-49b1-aabb-c54614aeadea"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First call:  4.0\n",
      "Saved global:  Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>\n"
     ]
    }
   ],
   "source": [
    "g = 0.\n",
    "def impure_saves_global(x):\n",
    "  global g\n",
    "  g = x\n",
    "  return x\n",
    "\n",
    "# JAX runs once the transformed function with special Traced values for arguments\n",
    "print (\"First call: \", jit(impure_saves_global)(4.))\n",
    "print (\"Saved global: \", g)  # Saved global has an internal JAX value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mlc2pQlp6v-9"
   },
   "source": [
    "A Python function can be functionally pure even if it actually uses stateful objects internally, as long as it does not read or write external state:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "TP-Mqf_862C0",
    "outputId": "78d55886-54de-483c-e7c4-bafd1d2c7219"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50.0\n"
     ]
    }
   ],
   "source": [
    "def pure_uses_internal_state(x):\n",
    "  state = dict(even=0, odd=0)\n",
    "  for i in range(10):\n",
    "    state['even' if i % 2 == 0 else 'odd'] += x\n",
    "  return state['even'] + state['odd']\n",
    "\n",
    "print(jit(pure_uses_internal_state)(5.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cDpQ5u63Ba_H"
   },
   "source": [
    "It is not recommended to use iterators in any JAX function you want to `jit` or in any control-flow primitive. The reason is that an iterator is a python object which introduces state to retrieve the next element. Therefore, it is incompatible with JAX's functional programming model. In the code below, there are some examples of incorrect attempts to use iterators with JAX. Most of them return an error, but some give unexpected results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "w99WXa6bBa_H",
    "outputId": "52d885fd-0239-4a08-f5ce-0c38cc008903"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "45\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "import jax.numpy as jnp\n",
    "from jax import make_jaxpr\n",
    "\n",
    "# lax.fori_loop\n",
    "array = jnp.arange(10)\n",
    "print(lax.fori_loop(0, 10, lambda i,x: x+array[i], 0)) # expected result 45\n",
    "iterator = iter(range(10))\n",
    "print(lax.fori_loop(0, 10, lambda i,x: x+next(iterator), 0)) # unexpected result 0\n",
    "\n",
    "# lax.scan\n",
    "def func11(arr, extra):\n",
    "    ones = jnp.ones(arr.shape)\n",
    "    def body(carry, aelems):\n",
    "        ae1, ae2 = aelems\n",
    "        return (carry + ae1 * ae2 + extra, carry)\n",
    "    return lax.scan(body, 0., (arr, ones))\n",
    "make_jaxpr(func11)(jnp.arange(16), 5.)\n",
    "# make_jaxpr(func11)(iter(range(16)), 5.) # throws error\n",
    "\n",
    "# lax.cond\n",
    "array_operand = jnp.array([0.])\n",
    "lax.cond(True, lambda x: x+1, lambda x: x-1, array_operand)\n",
    "iter_operand = iter(range(10))\n",
    "# lax.cond(True, lambda x: next(x)+1, lambda x: next(x)-1, iter_operand) # throws error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oBdKtkVW8Lha"
   },
   "source": [
    "## 🔪 In-place updates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JffAqnEW4JEb"
   },
   "source": [
    "In Numpy you're used to doing this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "om4xV7_84N9j",
    "outputId": "88b0074a-4440-41f6-caa7-031ac2d1a96f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "original array:\n",
      "[[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n",
      "updated array:\n",
      "[[0. 0. 0.]\n",
      " [1. 1. 1.]\n",
      " [0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "numpy_array = np.zeros((3,3), dtype=np.float32)\n",
    "print(\"original array:\")\n",
    "print(numpy_array)\n",
    "\n",
    "# In place, mutating update\n",
    "numpy_array[1, :] = 1.0\n",
    "print(\"updated array:\")\n",
    "print(numpy_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "go3L4x3w4-9p"
   },
   "source": [
    "If we try to do in-place indexed updating on a `jax.Array`, however, we get an __error__!  (☉_☉)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "iOscaa_GecEK",
    "outputId": "26fdb703-a476-4b7f-97ba-d28997ef750c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exception reporting mode: Minimal\n"
     ]
    }
   ],
   "source": [
    "%xmode Minimal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "2AxeCufq4wAp",
    "outputId": "fa4a87ad-1a84-471a-a3c5-a1396c432c85",
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m '<class 'jaxlib._jax.ArrayImpl'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html\n"
     ]
    }
   ],
   "source": [
    "jax_array = jnp.zeros((3,3), dtype=jnp.float32)\n",
    "\n",
    "# In place update of JAX's array will yield an error!\n",
    "jax_array[1, :] = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f520bec",
   "metadata": {},
   "source": [
    "And if we try to do `__iadd__`-style in-place updating, we get __different behavior than NumPy__!  (☉_☉)  (☉_☉)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20fbed45",
   "metadata": {},
   "outputs": [],
   "source": [
    "jax_array = jnp.array([10, 20])\n",
    "jax_array_new = jax_array\n",
    "jax_array_new += 10\n",
    "print(jax_array_new)  # `jax_array_new` is rebound to a new value [20, 30], but...\n",
    "print(jax_array)      # the original value is unmodified as [10, 20] !\n",
    "\n",
    "numpy_array = np.array([10, 20])\n",
    "numpy_array_new = numpy_array\n",
    "numpy_array_new += 10\n",
    "print(numpy_array_new)  # `numpy_array_new is numpy_array`, and it was updated\n",
    "print(numpy_array)      # in-place, so both are [20, 30] !"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2604e220",
   "metadata": {},
   "source": [
    "That's because NumPy defines `__iadd__` to perform in-place mutation. In\n",
    "contrast, `jax.Array` doesn't define an `__iadd__`, so Python treats\n",
    "`jax_array_new += 10` as syntactic sugar for `jax_array_new = jax_array_new +\n",
    "10`, rebinding the variable without mutating any arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7mo76sS25Wco"
   },
   "source": [
    "Allowing mutation of variables in-place makes program analysis and transformation difficult. JAX requires that programs are pure functions.\n",
    "\n",
    "Instead, JAX offers a _functional_ array update using the [`.at` property on JAX arrays](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hfloZ1QXCS_J"
   },
   "source": [
    "️⚠️ inside `jit`'d code and `lax.while_loop` or `lax.fori_loop` the __size__ of slices can't be functions of argument _values_ but only functions of argument _shapes_ -- the slice start indices have no such restriction.  See the below __Control Flow__ Section for more information on this limitation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X2Xjjvd-l8NL"
   },
   "source": [
    "### Array updates: `x.at[idx].set(y)`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SHLY52KQEiuX"
   },
   "source": [
    "For example, the update above can be written as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "PBGI-HIeCP_s",
    "outputId": "de13f19a-2066-4df1-d503-764c34585529"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "updated array:\n",
      " [[0. 0. 0.]\n",
      " [1. 1. 1.]\n",
      " [0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "jax_array = jnp.zeros((3,3), dtype=jnp.float32)\n",
    "updated_array = jax_array.at[1, :].set(1.0)\n",
    "print(\"updated array:\\n\", updated_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zUANAw9sCmgu"
   },
   "source": [
    "JAX's array update functions, unlike their NumPy versions, operate out-of-place. That is, the updated array is returned as a new array and the original array is not modified by the update."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "dbB0UmMhCe8f",
    "outputId": "55d46fa1-d0de-4c43-996c-f3bbc87b7175"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "original array unchanged:\n",
      " [[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n"
     ]
    }
   ],
   "source": [
    "print(\"original array unchanged:\\n\", jax_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eM6MyndXL2NY"
   },
   "source": [
    "However, inside __jit__-compiled code, if the __input value__ `x` of `x.at[idx].set(y)` is not reused, the compiler will optimize the array update to occur _in-place_."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7to-sF8EmC_y"
   },
   "source": [
    "### Array updates with other operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZY5l3tAdDmsJ"
   },
   "source": [
    "Indexed array updates are not limited simply to overwriting values. For example, we can perform indexed addition as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "tsw2svao8FUp",
    "outputId": "3c62a3b1-c12d-46f0-da74-791ec4b61e0b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "original array:\n",
      "[[1. 1. 1. 1. 1. 1.]\n",
      " [1. 1. 1. 1. 1. 1.]\n",
      " [1. 1. 1. 1. 1. 1.]\n",
      " [1. 1. 1. 1. 1. 1.]\n",
      " [1. 1. 1. 1. 1. 1.]]\n",
      "new array post-addition:\n",
      "[[1. 1. 1. 8. 8. 8.]\n",
      " [1. 1. 1. 1. 1. 1.]\n",
      " [1. 1. 1. 8. 8. 8.]\n",
      " [1. 1. 1. 1. 1. 1.]\n",
      " [1. 1. 1. 8. 8. 8.]]\n"
     ]
    }
   ],
   "source": [
    "print(\"original array:\")\n",
    "jax_array = jnp.ones((5, 6))\n",
    "print(jax_array)\n",
    "\n",
    "new_jax_array = jax_array.at[::2, 3:].add(7.)\n",
    "print(\"new array post-addition:\")\n",
    "print(new_jax_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sTjJ3WuaDyqU"
   },
   "source": [
    "For more details on indexed array updates, see the [documentation for the `.at` property](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oZ_jE2WAypdL"
   },
   "source": [
    "## 🔪 Out-of-bounds indexing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "btRFwEVzypdN"
   },
   "source": [
    "In Numpy, you are used to errors being thrown when you index an array outside of its bounds, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "5_ZM-BJUypdO",
    "outputId": "c9c41ae8-2653-4219-e6dc-09b03faa3b95",
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31mIndexError\u001b[0m\u001b[0;31m:\u001b[0m index 11 is out of bounds for axis 0 with size 10\n"
     ]
    }
   ],
   "source": [
    "np.arange(10)[11]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eoXrGARWypdR"
   },
   "source": [
    "However, raising an error from code running on an accelerator can be difficult or impossible. Therefore, JAX must choose some non-error behavior for out of bounds indexing (akin to how invalid floating point arithmetic results in `NaN`). When the indexing operation is an array index update (e.g. `index_add` or `scatter`-like primitives), updates at out-of-bounds indices will be skipped; when the operation is an array index retrieval (e.g. NumPy indexing or `gather`-like primitives) the index is clamped to the bounds of the array since __something__ must be returned. For example, the last value of the array will be returned from this indexing operation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "cusaAD0NypdR",
    "outputId": "af1708aa-b50b-4da8-f022-7f2fa67030a8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(9, dtype=int32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.arange(10)[11]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NAcXJNAcDi_v"
   },
   "source": [
    "If you would like finer-grained control over the behavior for out-of-bound indices, you can use the optional parameters of [`ndarray.at`](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html); for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "-0-MaFddO-xy",
    "outputId": "746c4b2b-a90e-4ec9-de56-ed6682d451e5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(9., dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.arange(10.0).at[11].get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "g5JEJtIUPBXi",
    "outputId": "4a0f6854-1165-47f2-e1ac-5a21fa2b8516"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(nan, dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.arange(10.0).at[11].get(mode='fill', fill_value=jnp.nan)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J8uO8yevBa_M"
   },
   "source": [
    "Note that due to this behavior for index retrieval, functions like `jnp.nanargmin` and `jnp.nanargmax` return -1 for slices consisting of NaNs whereas Numpy would throw an error.\n",
    "\n",
    "Note also that, as the two behaviors described above are not inverses of each other, reverse-mode automatic differentiation (which turns index updates into index retrievals and vice versa) [will not preserve the semantics of out of bounds indexing](https://github.com/jax-ml/jax/issues/5760). Thus it may be a good idea to think of out-of-bounds indexing in JAX as a case of [undefined behavior](https://en.wikipedia.org/wiki/Undefined_behavior)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LwB07Kx5sgHu"
   },
   "source": [
    "## 🔪 Non-array inputs: NumPy vs. JAX\n",
    "\n",
    "NumPy is generally happy accepting Python lists or tuples as inputs to its API functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "sErQES14sjCG",
    "outputId": "601485ff-4cda-48c5-f76c-2789073c4591"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum([1, 2, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZJ1Wt1bTtrSA"
   },
   "source": [
    "JAX departs from this, generally returning a helpful error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "DFEGcENSsmEc",
    "outputId": "08535679-6c1f-4dd9-a414-d8b59310d1ee",
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m sum requires ndarray or scalar arguments, got <class 'list'> at position 0.\n"
     ]
    }
   ],
   "source": [
    "jnp.sum([1, 2, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QPliLUZztxJt"
   },
   "source": [
    "This is a deliberate design choice, because passing lists or tuples to traced functions can lead to silent performance degradation that might otherwise be difficult to detect.\n",
    "\n",
    "For example, consider the following permissive version of `jnp.sum` that allows list inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "jhe-L_TwsvKd",
    "outputId": "ab2ee183-d9ec-45cc-d6be-5009347e1bc5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(45, dtype=int32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def permissive_sum(x):\n",
    "  return jnp.sum(jnp.array(x))\n",
    "\n",
    "x = list(range(10))\n",
    "permissive_sum(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m0XZLP7nuYdE"
   },
   "source": [
    "The output is what we would expect, but this hides potential performance issues under the hood. In JAX's tracing and JIT compilation model, each element in a Python list or tuple is treated as a separate JAX variable, and individually processed and pushed to device. This can be seen in the jaxpr for the ``permissive_sum`` function above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "k81u6DQ7vAjQ",
    "outputId": "869fc3b9-feda-4aa9-d2e5-5b5107de102d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ lambda ; a:i32[] b:i32[] c:i32[] d:i32[] e:i32[] f:i32[] g:i32[] h:i32[] i:i32[]\n",
       "    j:i32[]. let\n",
       "    k:i32[] = convert_element_type[new_dtype=int32 weak_type=False] a\n",
       "    l:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b\n",
       "    m:i32[] = convert_element_type[new_dtype=int32 weak_type=False] c\n",
       "    n:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d\n",
       "    o:i32[] = convert_element_type[new_dtype=int32 weak_type=False] e\n",
       "    p:i32[] = convert_element_type[new_dtype=int32 weak_type=False] f\n",
       "    q:i32[] = convert_element_type[new_dtype=int32 weak_type=False] g\n",
       "    r:i32[] = convert_element_type[new_dtype=int32 weak_type=False] h\n",
       "    s:i32[] = convert_element_type[new_dtype=int32 weak_type=False] i\n",
       "    t:i32[] = convert_element_type[new_dtype=int32 weak_type=False] j\n",
       "    u:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] k\n",
       "    v:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] l\n",
       "    w:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] m\n",
       "    x:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] n\n",
       "    y:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] o\n",
       "    z:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] p\n",
       "    ba:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] q\n",
       "    bb:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] r\n",
       "    bc:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] s\n",
       "    bd:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] t\n",
       "    be:i32[10] = concatenate[dimension=0] u v w x y z ba bb bc bd\n",
       "    bf:i32[] = reduce_sum[axes=(0,)] be\n",
       "  in (bf,) }"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "make_jaxpr(permissive_sum)(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C0_dpCfpvCts"
   },
   "source": [
    "Each entry of the list is handled as a separate input, resulting in a tracing & compilation overhead that grows linearly with the size of the list. To prevent surprises like this, JAX avoids implicit conversions of lists and tuples to arrays.\n",
    "\n",
    "If you would like to pass a tuple or list to a JAX function, you can do so by first explicitly converting it to an array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "nFf_DydixG8v",
    "outputId": "e31b43b3-05f7-4300-fdd2-40e3896f6f8f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(45, dtype=int32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.sum(jnp.array(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MUycRNh6e50W"
   },
   "source": [
    "## 🔪 Random numbers\n",
    "\n",
    "JAX's pseudo-random number generation differs from Numpy's in important ways. For a quick how-to, see {ref}`key-concepts-prngs`. For more details, see the {ref}`pseudorandom-numbers` tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dc0e6b2",
   "metadata": {
    "id": "rg4CpMZ8c3ri"
   },
   "source": [
    "## 🔪 Control flow\n",
    "\n",
    "Moved to {ref}`control-flow`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OxLsZUyRt_kF"
   },
   "source": [
    "## 🔪 Dynamic shapes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1tKXcAMduDR1"
   },
   "source": [
    "JAX code used within transforms like `jax.jit`, `jax.vmap`, `jax.grad`, etc. requires all output arrays and intermediate arrays to have static shape: that is, the shape cannot depend on values within other arrays.\n",
    "\n",
    "For example, if you were implementing your own version of `jnp.nansum`, you might start with something like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "9GIwgvfLujiD"
   },
   "outputs": [],
   "source": [
    "def nansum(x):\n",
    "  mask = ~jnp.isnan(x)  # boolean mask selecting non-nan values\n",
    "  x_without_nans = x[mask]\n",
    "  return x_without_nans.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "43S7wYAiupGe"
   },
   "source": [
    "Outside JIT and other transforms, this works as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "ITYoNQEZur4s",
    "outputId": "a9a03d25-9c54-43b6-d35e-aea6c448d680"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.0\n"
     ]
    }
   ],
   "source": [
    "x = jnp.array([1, 2, jnp.nan, 3, 4])\n",
    "print(nansum(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "guup5n8xvGI-"
   },
   "source": [
    "If you attempt to apply `jax.jit` or another transform to this function, it will error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "id": "nms9KjQEvNTz",
    "outputId": "d8ae982f-111d-45b6-99f8-37715e2eaab3",
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "NonConcreteBooleanIndexError",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31mNonConcreteBooleanIndexError\u001b[0m\u001b[0;31m:\u001b[0m Array boolean indices must be concrete; got ShapedArray(bool[5])\n\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError\n"
     ]
    }
   ],
   "source": [
    "jax.jit(nansum)(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r2aGyHDkvauu"
   },
   "source": [
    "The problem is that the size of `x_without_nans` is dependent on the values within `x`, which is another way of saying its size is *dynamic*.\n",
    "Often in JAX it is possible to work-around the need for dynamically-sized arrays via other means.\n",
    "For example, here it is possible to use the three-argument form of  `jnp.where` to replace the NaN values with zeros, thus computing the same result while avoiding dynamic shapes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "id": "Zbuj7Dg1wnSg",
    "outputId": "81a5e356-cd28-4709-b307-07c6254c82de"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.0\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def nansum_2(x):\n",
    "  mask = ~jnp.isnan(x)  # boolean mask selecting non-nan values\n",
    "  return jnp.where(mask, x, 0).sum()\n",
    "\n",
    "print(nansum_2(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uGH-jqK7wxTl"
   },
   "source": [
    "Similar tricks can be played in other situations where dynamically-shaped arrays occur."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DKTMw6tRZyK2"
   },
   "source": [
    "## 🔪 Debugging NaNs and Infs\n",
    "\n",
    "Use the `jax_debug_nans` and `jax_debug_infs` flags to find the source of NaN/Inf values in functions and gradients. See {ref}`debugging-flags`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YTktlwTTMgFl"
   },
   "source": [
    "## 🔪 Double (64bit) precision\n",
    "\n",
    "At the moment, JAX by default enforces single-precision numbers to mitigate the Numpy API's tendency to aggressively promote operands to `double`.  This is the desired behavior for many machine-learning applications, but it may catch you by surprise!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "id": "CNNGtzM3NDkO",
    "outputId": "b422bb23-a784-44dc-f8c9-57f3b6c861b8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dtype('float32')"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = random.uniform(random.key(0), (1000,), dtype=jnp.float64)\n",
    "x.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VcvqzobxNPbd"
   },
   "source": [
    "To use double-precision numbers, you need to set the `jax_enable_x64` configuration variable __at startup__.\n",
    "\n",
    "There are a few ways to do this:\n",
    "\n",
    "1. You can enable 64-bit mode by setting the environment variable `JAX_ENABLE_X64=True`.\n",
    "\n",
    "2. You can manually set the `jax_enable_x64` configuration flag at startup:\n",
    "\n",
    "   ```python\n",
    "   # again, this only works on startup!\n",
    "   import jax\n",
    "   jax.config.update(\"jax_enable_x64\", True)\n",
    "   ```\n",
    "\n",
    "3. You can parse command-line flags with `absl.app.run(main)`\n",
    "\n",
    "   ```python\n",
    "   import jax\n",
    "   jax.config.config_with_absl()\n",
    "   ```\n",
    "\n",
    "4. If you want JAX to run absl parsing for you, i.e. you don't want to do `absl.app.run(main)`, you can instead use\n",
    "\n",
    "   ```python\n",
    "   import jax\n",
    "   if __name__ == '__main__':\n",
    "     # calls jax.config.config_with_absl() *and* runs absl parsing\n",
    "     jax.config.parse_flags_with_absl()\n",
    "   ```\n",
    "\n",
    "Note that #2-#4 work for _any_ of JAX's configuration options.\n",
    "\n",
    "We can then confirm that `x64` mode is enabled, for example:\n",
    "\n",
    "```python\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import random\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)\n",
    "x = random.uniform(random.key(0), (1000,), dtype=jnp.float64)\n",
    "x.dtype # --> dtype('float64')\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6Cks2_gKsXaW"
   },
   "source": [
    "### Caveats\n",
    "⚠️ XLA doesn't support 64-bit convolutions on all backends!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WAHjmL0E2XwO"
   },
   "source": [
    "## 🔪 Miscellaneous divergences from NumPy\n",
    "\n",
    "While `jax.numpy` makes every attempt to replicate the behavior of numpy's API, there do exist corner cases where the behaviors differ.\n",
    "Many such cases are discussed in detail in the sections above; here we list several other known places where the APIs diverge.\n",
    "\n",
    "- For binary operations, JAX's type promotion rules differ somewhat from those used by NumPy. See [Type Promotion Semantics](https://docs.jax.dev/en/latest/type_promotion.html) for more details.\n",
    "- When performing unsafe type casts (i.e. casts in which the target dtype cannot represent the input value), JAX's behavior may be backend dependent, and in general may diverge from NumPy's behavior. Numpy allows control over the result in these scenarios via the `casting` argument (see [`np.ndarray.astype`](https://numpy.org/devdocs/reference/generated/numpy.ndarray.astype.html)); JAX does not provide any such configuration, instead directly inheriting the behavior of [XLA:ConvertElementType](https://www.openxla.org/xla/operation_semantics#convertelementtype).\n",
    "\n",
    "  Here is an example of an unsafe cast with differing results between NumPy and JAX:\n",
    "  ```python\n",
    "  >>> np.arange(254.0, 258.0).astype('uint8')\n",
    "  array([254, 255,   0,   1], dtype=uint8)\n",
    "\n",
    "  >>> jnp.arange(254.0, 258.0).astype('uint8')\n",
    "  Array([254, 255, 255, 255], dtype=uint8)\n",
    "\n",
    "  ```\n",
    "  This sort of mismatch would typically arise when casting extreme values from floating to integer types or vice versa.\n",
    "- When operating on [subnormal](https://en.wikipedia.org/wiki/Subnormal_number)\n",
    "  floating point numbers, JAX operations use flush-to-zero semantics on some\n",
    "  backends. For example:\n",
    "  ```python\n",
    "  >>> import jax.numpy as jnp\n",
    "  >>> subnormal = jnp.float32(1E-45)\n",
    "  >>> subnormal  # subnormals are representable\n",
    "  Array(1.e-45, dtype=float32)\n",
    "  >>> subnormal + 0  # but are flushed to zero within operations\n",
    "  Array(0., dtype=float32)\n",
    "\n",
    "  ```\n",
    "  The detailed operation semantics for subnormal values will generally\n",
    "  vary depending on the backend.\n",
    "\n",
    "## 🔪 Sharp bits covered in tutorials\n",
    "- {ref}`control-flow` discusses how to work with the constraints that `jit` imposes on the use of Python control flow and logical operators.\n",
    "- {ref}`stateful-computations` gives some advice on how to properly handle state in a JAX program, given that JAX transformations can be applied only to pure functions.\n",
    "\n",
    "## Fin.\n",
    "\n",
    "If something's not covered here that has caused you weeping and gnashing of teeth, please let us know and we'll extend these introductory _advisos_!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Common Gotchas in JAX",
   "provenance": [],
   "toc_visible": true
  },
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2 (v3.8.2:7b3ab5921f, Feb 24 2020, 17:52:18) \n[Clang 6.0 (clang-600.0.57)]"
  },
  "mystnb": {
   "render_error_lexer": "none"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
